#!/usr/bin/env python
# coding=utf-8

import os
import sys

import kaldi_io
import numpy as np
import argparse
from sklearn import mixture
import joblib




if __name__ == "__main__":
	parser = argparse.ArgumentParser()
	parser.add_argument("--npy-path", type=str, default="dvector.npy", help="npy file path")
	parser.add_argument("--out-dir", type=str, default="out/", help="numpy out dir")
	parser.add_argument("--n-components", type=int, default=32, help="gmm n_components")
	parser.add_argument("--max-iter", type=int, default=8, help="gmm n_components")
	parser.add_argument("--covariance-type", type=str, default="diag", \
			help="GMM covariance_type: spherical, tied, diag, full")
	args = parser.parse_args()

	dst_path = os.path.join(args.out_dir, "gmm_{}.{}".format(args.n_components, args.covariance_type))
	if os.path.exists(dst_path):
		print("pass: {}".format(dst_path))
		sys.exit(0)
	else:
		data = np.load(args.npy_path)
		print("{}\tshape: {}\ttype: {}".format(args.npy_path.split("/")[-2], data.shape, args.covariance_type))
		gmm = mixture.GaussianMixture(n_components=args.n_components, reg_covar=1e-02, \
				covariance_type=args.covariance_type, max_iter=args.max_iter)
		gmm.fit(data)
		joblib.dump(gmm, dst_path)

